#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Gravity calibration & online compensation (ROS2, rclpy)
同步补偿：/wrench_comp 发布频率与 /force_data 保持一致
重力补偿实验用配合force_three6.py使用
"""
import os
import json
from pathlib import Path
import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.duration import Duration
from rclpy.time import Time

import tf2_ros
from geometry_msgs.msg import WrenchStamped
from std_srvs.srv import Trigger


def quat_to_rotmat(x, y, z, w):
    n = np.sqrt(x * x + y * y + z * z + w * w)
    if n == 0.0:
        return np.eye(3, dtype=float)
    x, y, z, w = x / n, y / n, z / n, w / n
    xx, yy, zz = x * x, y * y, z * z
    xy, xz, yz = x * y, x * z, y * z
    wx, wy, wz = w * x, w * y, w * z
    return np.array([
        [1 - 2 * (yy + zz), 2 * (xy - wz),     2 * (xz + wy)],
        [2 * (xy + wz),     1 - 2 * (xx + zz), 2 * (yz - wx)],
        [2 * (xz - wy),     2 * (yz + wx),     1 - 2 * (xx + yy)]
    ], dtype=float)


CALIB_FILE = Path.home() / ".ros" / "gravity_calib.json"


class GravityCalibNode(Node):
    def __init__(self):
        super().__init__('gravity_calib_node')

        # ---------------- 参数 ----------------
        self.declare_parameter('base_frame', 'base_link')
        self.declare_parameter('ee_frame', 'end')
        self.declare_parameter('wrench_topic', '/force_data')
        self.declare_parameter('capture_timeout_s', 2.0)
        self.declare_parameter('g_B', None)
        self.declare_parameter('F0', None)

        self.declare_parameter('enable_online_bias_tune', True)
        self.declare_parameter('no_contact_thresh', 1.0)
        self.declare_parameter('alpha_bias', 0.001)

        self.base_frame = self.get_parameter('base_frame').value
        self.ee_frame = self.get_parameter('ee_frame').value
        self.wrench_topic = self.get_parameter('wrench_topic').value
        self.cap_timeout = float(self.get_parameter('capture_timeout_s').value)

        self.enable_online_bias_tune = bool(self.get_parameter('enable_online_bias_tune').value)
        self.no_contact_thresh = float(self.get_parameter('no_contact_thresh').value)
        self.alpha_bias = float(self.get_parameter('alpha_bias').value)

        # 末端->力传感器固定旋转
        self.M = np.array([[0, 0, 1],
                           [0, -1, 0],
                           [1, 0, 0]], dtype=float)

        # ---------------- TF / Topics ----------------
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)

        self.last_force = None
        self.sub_wrench = self.create_subscription(
            WrenchStamped, self.wrench_topic, self.cb_wrench, 10
        )
        self.pub_comp = self.create_publisher(WrenchStamped, '/wrench_comp', 10)

        # ---------------- 标定与样本 ----------------
        self.samples_RSB = []
        self.samples_F = []

        self.g_B = None
        self.F0 = None

        self._load_calib_from_params_or_file()

        # ---------------- 服务 ----------------
        self.srv_capture = self.create_service(Trigger, 'capture_sample', self.handle_capture)
        self.srv_solve = self.create_service(Trigger, 'solve_calibration', self.handle_solve)
        self.srv_save = self.create_service(Trigger, 'save_calibration', self.handle_save)
        self.srv_clear = self.create_service(Trigger, 'clear_calibration', self.handle_clear)
        self.srv_reload = self.create_service(Trigger, 'reload_calibration', self.handle_reload)

        self.get_logger().info(
            f'GravityCalibNode up. base={self.base_frame}, ee={self.ee_frame}, '
            f'wrench={self.wrench_topic}, calib_file={str(CALIB_FILE)}'
        )

    # ===== 基础回调 =====
    def cb_wrench(self, msg: WrenchStamped):
        self.last_force = np.array([msg.wrench.force.x,
                                    msg.wrench.force.y,
                                    msg.wrench.force.z], dtype=float)
        # 收到新数据时立即做补偿并发布
        self.online_compensation(msg)

    def lookup_R_SB(self) -> np.ndarray:
        now = Time()
        ok = self.tf_buffer.can_transform(self.base_frame, self.ee_frame, now,
                                          timeout=Duration(seconds=self.cap_timeout))
        if not ok:
            raise RuntimeError(f'等待 TF 超时: {self.base_frame}->{self.ee_frame}')

        trans = self.tf_buffer.lookup_transform(self.base_frame, self.ee_frame, now,
                                                timeout=Duration(seconds=0.5))
        q = trans.transform.rotation
        R_BE = quat_to_rotmat(q.x, q.y, q.z, q.w)
        R_BS = R_BE @ self.M
        R_SB = R_BS.T
        return R_SB

    # ===== 标定采样/求解 =====
    def handle_capture(self, request, response):
        try:
            if self.last_force is None:
                raise RuntimeError('暂未收到力传感器数据')
            R_SB = self.lookup_R_SB()
            self.samples_RSB.append(R_SB)
            self.samples_F.append(self.last_force.copy())
            n = len(self.samples_F)
            response.success = True
            response.message = f'采样成功，第 {n} 组。'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'采样失败: {e}'
            self.get_logger().warn(response.message)
        return response

    def handle_solve(self, request, response):
        try:
            N = len(self.samples_F)
            if N < 3:
                raise RuntimeError(f'样本不足，需要≥3，当前 {N}')

            A_rows, b_rows = [], []
            for i in range(N):
                R = self.samples_RSB[i]
                F = self.samples_F[i]
                A_rows.append(np.hstack([R, np.eye(3)]))
                b_rows.append(F)
            A = np.vstack(A_rows)
            b = np.hstack(b_rows)

            x, *_ = np.linalg.lstsq(A, b, rcond=None)
            self.g_B = x[0:3]
            self.F0 = x[3:6]

            g_norm = float(np.linalg.norm(self.g_B))
            self.get_logger().info(f'求解完成: g_B={self.g_B} (|g_B|={g_norm:.4f}), F0={self.F0}')

            self._save_calib_to_file()

            response.success = True
            response.message = '标定完成并已保存，进入在线补偿模式。'
        except Exception as e:
            response.success = False
            response.message = f'求解失败: {e}'
            self.get_logger().error(response.message)
        return response

    # ===== 文件持久化与维护服务 =====
    def _validate_vec(self, v):
        return isinstance(v, (list, tuple)) and len(v) == 3

    def _load_calib_from_params_or_file(self):
        g_param = self.get_parameter('g_B').value
        F_param = self.get_parameter('F0').value
        if self._validate_vec(g_param) and self._validate_vec(F_param):
            self.g_B = np.array(g_param, dtype=float)
            self.F0 = np.array(F_param, dtype=float)
            self.get_logger().info(f'从参数加载标定：g_B={self.g_B}, F0={self.F0}')
            return

        if CALIB_FILE.exists():
            try:
                with open(CALIB_FILE, 'r') as f:
                    data = json.load(f)
                if (self._validate_vec(data.get('g_B')) and
                        self._validate_vec(data.get('F0')) and
                        data.get('base_frame') == self.base_frame and
                        data.get('ee_frame') == self.ee_frame):
                    self.g_B = np.array(data['g_B'], dtype=float)
                    self.F0 = np.array(data['F0'], dtype=float)
                    self.get_logger().info(f'从文件加载标定：g_B={self.g_B}, F0={self.F0}')
                else:
                    self.get_logger().warn('标定文件存在，但与当前配置不匹配或字段缺失，忽略。')
            except Exception as e:
                self.get_logger().warn(f'读取标定文件失败：{e}')

    def _save_calib_to_file(self):
        if self.g_B is None or self.F0 is None:
            raise RuntimeError('当前无标定数据可保存')
        CALIB_FILE.parent.mkdir(parents=True, exist_ok=True)
        data = {
            'g_B': self.g_B.tolist(),
            'F0': self.F0.tolist(),
            'base_frame': self.base_frame,
            'ee_frame': self.ee_frame,
            'M': self.M.tolist()
        }
        with open(CALIB_FILE, 'w') as f:
            json.dump(data, f, indent=2)
        self.get_logger().info(f'标定已保存：{CALIB_FILE}')

    def handle_save(self, request, response):
        try:
            self._save_calib_to_file()
            response.success = True
            response.message = '已保存到文件。'
        except Exception as e:
            response.success = False
            response.message = f'保存失败: {e}'
        return response

    def handle_clear(self, request, response):
        try:
            if CALIB_FILE.exists():
                CALIB_FILE.unlink()
            self.g_B, self.F0 = None, None
            response.success = True
            response.message = '已清除本地标定与内存。'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'清除失败: {e}'
            self.get_logger().error(response.message)
        return response

    def handle_reload(self, request, response):
        try:
            self.g_B, self.F0 = None, None
            self._load_calib_from_params_or_file()
            ok = (self.g_B is not None and self.F0 is not None)
            response.success = ok
            response.message = '已重载标定。' if ok else '未找到可用标定。'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'重载失败: {e}'
            self.get_logger().error(response.message)
        return response

    # ===== 在线补偿（订阅触发）=====
    def online_compensation(self, msg_in: WrenchStamped):
        if self.g_B is None or self.F0 is None:
            return
        try:
            R_SB = self.lookup_R_SB()
            F_pred = R_SB @ self.g_B + self.F0
            F_comp = self.last_force - F_pred

            if self.enable_online_bias_tune:
                if np.linalg.norm(F_comp) < self.no_contact_thresh:
                    self.F0 = (1.0 - self.alpha_bias) * self.F0 + \
                              self.alpha_bias * (self.last_force - R_SB @ self.g_B)

            msg_out = WrenchStamped()
            msg_out.header.stamp = msg_in.header.stamp  # 使用 force_data 时间戳
            msg_out.header.frame_id = 'sensor'
            msg_out.wrench.force.x = -float(F_comp[0])/10
            msg_out.wrench.force.y = float(F_comp[1])/40
            msg_out.wrench.force.z = -float(F_comp[2])/10
            self.pub_comp.publish(msg_out)

            self.get_logger().info(f"wrench_comp:: Fx={msg_out.wrench.force.x:.3f}, Fy1={msg_out.wrench.force.y:.3f}， Fz={msg_out.wrench.force.z:.3f}")
        except Exception as e:
            self.get_logger().warn(f'在线补偿失败: {e}')


def main():
    rclpy.init()
    node = GravityCalibNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
